//go:build windows
// +build windows

package win

import (
	"dgo/goutils/utils"
	"dgo/goutils/ylibos"
	"fmt"
	"os/exec"
	"syscall"
	"unsafe"
)

type PROCESSENTRY32 struct {
	dwSize              ulong     // 结构大小
	cntUsage            ulong     // 此进程的引用计数
	th32ProcessID       ulong     // 进程id
	th32DefaultHeapID   ulong_ptr // 进程默认堆id
	th32ModuleID        ulong     // 进程模块id
	cntThreads          ulong     // 进程的线程数
	th32ParentProcessID ulong     // 父进程id
	pcPriClassBase      ulong     // 线程优先权
	dwFlags             ulong     // 保留
	szExeFile           [260]byte // 进程全名
}

func (this *PROCESSENTRY32) GetPName() string {
	return utils.Trim(string(utils.TrimRightBytes(this.szExeFile[0:], 0)))
}

func (this *PROCESSENTRY32) GetPID() ulong {
	return this.th32ProcessID
}

func (this *PROCESSENTRY32) GetModuleID() ulong {
	return this.th32ModuleID
}

type MODULEENTRY32 struct {
	dwSize        ulong     // 指定结构的长度，以字节为单位。在调用Module32First功能，设置这个成员SIZEOF（MODULEENTRY32）。如果你不初始化的dwSize，Module32First将失败。
	th32ModuleID  ulong     // 此成员已经不再被使用，通常被设置为1
	th32ProcessID ulong     // 正在检查的进程标识符。这个成员的内容，可以使用Win32 API的元素
	GlblcntUsage  ulong     // 全局模块的使用计数，即模块的总载入次数。通常这一项是没有意义的，被设置为0xFFFF。
	ProccntUsage  ulong     // 全局模块的使用计数（与GlblcntUsage相同）。通常这一项也是没有意义的，被设置为0xFFFF。
	modBaseAddr   byte      // 模块的基址，在其所属的进程范围内。
	modBaseSize   ulong     // 模块的大小，单位字节。
	hModule       ulong_ptr // 模块句柄
	szModule      [260]byte // 模块名称
	szExePath     [260]byte // NULL结尾的字符串，其中包含的位置，或模块的路径。 在VC++6.0中， _MAX_PATH的值为260。
}

func RangeProcessList(fn func(info *ylibos.ProcessInfo) bool) error {
	pHandle, _, _ := CreateToolhelp32Snapshot.Call(uintptr(0x2), uintptr(0x0))
	if int(pHandle) == -1 {
		return fmt.Errorf("CreateToolhelp32Snapshot err")
	}
	defer CloseHandle.Call(pHandle)

	for {
		var info ylibos.ProcessInfo
		var proc PROCESSENTRY32
		proc.dwSize = ulong(unsafe.Sizeof(proc))
		if rt, _, _ := Process32Next.Call(uintptr(pHandle), uintptr(unsafe.Pointer(&proc))); int(rt) == 1 {
			info.PID = int(proc.th32ProcessID)
			info.PName = proc.GetPName()
			e1 := ReadProcessFullName2(uint32(proc.th32ProcessID), &info)
			if e1 != nil {
				info.LastErr = e1
			}
			if !fn(&info) {
				break
			}
		} else {
			break
		}
	}

	return nil
}

func ProcessKill(pid int) error {
	cmd := exec.Command("taskkill.exe", "/PID", fmt.Sprintf("%d", pid))
	err := cmd.Start()
	if err != nil {
		return err
	}
	go cmd.Wait()
	return nil
}

func ProcessKillWaitFor(pid int) error {
	cmd := exec.Command("taskkill.exe", "/PID", fmt.Sprintf("%d", pid))
	return cmd.Run()
}

func GetProcessList() []ylibos.ProcessInfo {
	var processArray []ylibos.ProcessInfo

	//pHandle, _, _ := CreateToolhelp32Snapshot.Call(uintptr(0x2), uintptr(0x0))

	pHandle, _, _ := CreateToolhelp32Snapshot.Call(uintptr(0x2), uintptr(0x0))
	if int(pHandle) == -1 {
		return processArray
	}

	for {
		var info ylibos.ProcessInfo
		var proc PROCESSENTRY32
		proc.dwSize = ulong(unsafe.Sizeof(proc))
		if rt, _, _ := Process32Next.Call(uintptr(pHandle), uintptr(unsafe.Pointer(&proc))); int(rt) == 1 {
			// fmt.Println("ProcessName : " + string(proc.szExeFile[0:]))
			// fmt.Println("th32ModuleID : " + strconv.Itoa(int(proc.th32ModuleID)))
			// fmt.Println("ProcessID : " + strconv.Itoa(int(proc.th32ProcessID)))
			info.PID = int(proc.th32ProcessID)
			info.PName = proc.GetPName()
			ReadProcessFullName2(uint32(proc.th32ProcessID), &info)
			processArray = append(processArray, info)
		} else {
			break
		}
	}
	_, _, _ = CloseHandle.Call(pHandle)

	return processArray
}

func GetProcess(pid int) (PROCESSENTRY32, error) {
	var targetProcess PROCESSENTRY32
	targetProcess = PROCESSENTRY32{
		dwSize: 0,
	}

	pHandle, _, _ := CreateToolhelp32Snapshot.Call(uintptr(0x2), uintptr(0x0))
	if int(pHandle) == -1 {
		return targetProcess, fmt.Errorf("error:Can not find any proess.")
	}
	Process32Next := kernel32.NewProc("Process32Next")

	for {
		var proc PROCESSENTRY32
		proc.dwSize = ulong(unsafe.Sizeof(proc))
		if rt, _, _ := Process32Next.Call(uintptr(pHandle), uintptr(unsafe.Pointer(&proc))); int(rt) == 1 {
			if int(proc.th32ProcessID) == pid {
				targetProcess = proc
				// fmt.Println("ProcessName : " + string(proc.szExeFile[0:]))
				// fmt.Println("th32ModuleID : " + strconv.Itoa(int(proc.th32ModuleID)))
				// fmt.Println("ProcessID : " + strconv.Itoa(int(proc.th32ProcessID)))
				break
			}
		} else {
			break
		}
	}

	_, _, _ = CloseHandle.Call(pHandle)
	return targetProcess, nil
}

type ListMode uintptr

const (
	LIST_MODULES_DEFAULT ListMode = 0x0  //查看默认的模块
	LIST_MODULES_32BIT            = 0x01 //查看32-bit的模块
	LIST_MODULES_64BIT            = 0x02 //查看64-bit的模块
	LIST_MODULES_ALL              = 0x03 //查看所有的模块
)

func ListDynamicModule(pid uint32, mode ListMode) ([]string, error) {
	//获取进程句柄
	var hProcess, err = syscall.OpenProcess(0x0400|0x0010, false, pid)
	if err != nil {
		return nil, err
	}

	//注意关闭句柄
	defer syscall.CloseHandle(hProcess)

	var (
		cbNeeded uint32
		hMods    [1024]uintptr
	)

	//获取进程打开的模块数
	r, _, err := enumProcessModulesEx.Call(uintptr(hProcess), uintptr(unsafe.Pointer(&hMods[0])), unsafe.Sizeof(hMods), uintptr(unsafe.Pointer(&cbNeeded)), uintptr(mode))
	if r == 0 {
		return nil, err
	}

	//计算模块数量
	var modelength = int(cbNeeded) / int(unsafe.Sizeof(hMods[0]))
	var szModName [260]uint16
	var list = make([]string, 0, modelength)

	for i := 0; i < modelength; i++ {
		r, _, err = getModuleFileNameExW.Call(uintptr(hProcess), hMods[i], uintptr(unsafe.Pointer(&szModName)), unsafe.Sizeof(szModName))
		if r == 0 {
			return list, err
		}
		list = append(list, syscall.UTF16ToString(szModName[:]))
	}
	return list, nil
}

func ReadProcessFullName(handle syscall.Handle, info *ylibos.ProcessInfo) error {
	var hMod uintptr
	var szModName [260]uint16

	r, _, err := getModuleFileNameExW.Call(uintptr(handle), hMod, uintptr(unsafe.Pointer(&szModName)), unsafe.Sizeof(szModName))
	if r == 0 {
		return err
	}
	info.ExeFileName = syscall.UTF16ToString(szModName[:])
	return nil
}

func ReadProcessFullName2(pid uint32, info *ylibos.ProcessInfo) error {
	var hProcess, err = syscall.OpenProcess(0x0400|0x0010, false, pid)
	if err != nil {
		return err
	}
	defer syscall.CloseHandle(hProcess)
	return ReadProcessFullName(hProcess, info)
	//return nil;
}

func GetProcessInfo(pid uint32) (*ylibos.ProcessInfo, error) {
	rval := &ylibos.ProcessInfo{}
	ReadProcessFullName2(pid, rval)
	return rval, nil
}
